package cn.xaut.优先队列;

import java.util.*;

/**
 * 347. 前 K 个高频元素
 */
public class demo347 {

    public int[] topKFrequent(int[] nums, int k) {

        Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            if (map.containsKey(num))
                map.put(num, map.get(num) + 1);
            else
                map.put(num, 1);
        }

        List<Map.Entry<Integer, Integer>> list = new ArrayList<>(map.entrySet());
        list.sort(new Comparator<Map.Entry<Integer, Integer>>() {
            @Override
            public int compare(Map.Entry<Integer, Integer> o1, Map.Entry<Integer, Integer> o2) {
                return o2.getValue() - o1.getValue();
            }
        });

        int[] res = new int[k];
        for (int i = 0; i < k; i++)
            res[i] = list.get(i).getKey();

        return res;
    }

    // 优先队列
    public int[] topKFrequent2(int[] nums, int k) {
        
        Map<Integer, Integer> map = new HashMap<>();
        for (int num : nums) {
            if (map.containsKey(num))
                map.put(num, map.get(num) + 1);
            else
                map.put(num, 1);
        }

        PriorityQueue<int[]> pq = new PriorityQueue<int[]>(new Comparator<int[]>() {
            @Override
            public int compare(int[] o1, int[] o2) {
                return o1[1] - o2[1];
            }
        });

        for (Map.Entry<Integer, Integer> entry : map.entrySet()) {
            int key = entry.getKey();
            int value = entry.getValue();
            if (pq.size() == k) {
                assert pq.peek() != null;
                if (value > pq.peek()[1]) {
                    pq.poll();
                    pq.offer(new int[]{key, value});
                }
            } else
                pq.offer(new int[]{key, value});
        }

        int[] res = new int[k];
        for (int i = 0; i < k; i++) 
            res[i] = Objects.requireNonNull(pq.poll())[0];  // res[i] = pq.poll()[0];
        
        return res;
    }

    public static void main(String[] args) {

        System.out.println(Arrays.toString(new demo347().topKFrequent2(new int[]{1, 1, 1, 2, 2, 3}, 2)));//[1, 2]
        System.out.println(Arrays.toString(new demo347().topKFrequent2(new int[]{1}, 1)));//[1]
    }
}
